from book import ContentType
from tenacity import retry, wait_random_exponential, stop_after_attempt
from utils import LOG

class Model:
    def __init__(self, model_name: str, api_key: str, client):
        self.model_name = model_name
        self.client = client(api_key=api_key)

    def make_text_prompt(self, text: str, target_language: str, delimiter='$$$') -> str:
        return f"""
            下面我让你来充当翻译家，你的目标是把任何语言翻译成{target_language}，
            请翻译时不要带翻译腔，而是要翻译得自然、流畅和地道，使用优美和高雅的表达方式。
            请翻译由{delimiter}分隔的内容：{delimiter} {text} {delimiter}
        """

    def make_table_prompt(self, table: str, target_language: str, delimiter='$$$') -> str:
        return f"""
            下面我让你来充当翻译家，你的目标是把任何语言翻译成{target_language}，
            请翻译时不要带翻译腔，而是要翻译得自然、流畅和地道，使用优美和高雅的表达方式。
            保持间距（空格，分隔符），必须以 Markdown 的表格语法形式返回，注意不要包含<```>。
            请翻译由{delimiter}分隔的内容：{delimiter} {table} {delimiter}
        """

    def translate_prompt(self, content, target_language: str) -> str:
        if content.content_type == ContentType.TEXT:
            return self.make_text_prompt(content.original, target_language)
        elif content.content_type == ContentType.TABLE:
            return self.make_table_prompt(content.original, target_language)

    @retry(wait=wait_random_exponential(multiplier=1, max=40), stop=stop_after_attempt(3))
    def make_request(self, prompt):
        try:
            response = self.client.chat.completions.create(
                messages=[
                    {"role": "user", "content": prompt},
                ],
                model=self.model_name,
            )
            return response.choices[0].message.content, True
        except Exception as e:
            LOG.error(f'Unable to generate ChatCompletion response.\nException: {e}')
            return None, False